package com.commonEncryptionAlgorithms.dissymmetry.rsa;

import javax.crypto.Cipher;
import java.io.*;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.security.*;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.util.Base64;

public class RSAUtil {
    private static final String ALGORITHM = "RSA";
    private static final int KEY_SIZE = 2048;

    /**
     * 生成RSA密钥对
     */
    private static KeyPair generateKeyPair() throws NoSuchAlgorithmException {
        KeyPairGenerator keyGen = KeyPairGenerator.getInstance(ALGORITHM);
        keyGen.initialize(KEY_SIZE);
        return keyGen.generateKeyPair();
    }

    /**
     * 将公钥导出为PEM格式的字符串
     */
    private static String exportPublicKeyToPEM(PublicKey publicKey) {
        byte[] encodedPublicKey = publicKey.getEncoded();
        return "-----BEGIN PUBLIC KEY-----\n" +
                Base64.getMimeEncoder(64, new byte[]{'\n'}).encodeToString(encodedPublicKey) +
                "\n-----END PUBLIC KEY-----";
    }

    /**
     * 将私钥导出为PEM格式的字符串
     */
    private static String exportPrivateKeyToPEM(PrivateKey privateKey) {
        byte[] encodedPrivateKey = privateKey.getEncoded();
        return "-----BEGIN PRIVATE KEY-----\n" +
                Base64.getMimeEncoder(64, new byte[]{'\n'}).encodeToString(encodedPrivateKey) +
                "\n-----END PRIVATE KEY-----";
    }

    /**
     * 从PEM格式的字符串导入公钥
     */
    private static PublicKey importPublicKeyFromPEM(String pemString) throws Exception {
        String publicKeyPEM = pemString.replace("-----BEGIN PUBLIC KEY-----", "")
                .replace("-----END PUBLIC KEY-----", "")
                .replaceAll("\\s+", "");
        byte[] encodedPublicKey = Base64.getMimeDecoder().decode(publicKeyPEM);
        X509EncodedKeySpec keySpec = new X509EncodedKeySpec(encodedPublicKey);
        KeyFactory keyFactory = KeyFactory.getInstance(ALGORITHM);
        return keyFactory.generatePublic(keySpec);
    }

    /**
     * 从PEM格式的字符串导入私钥
     */
    private static PrivateKey importPrivateKeyFromPEM(String pemString) throws Exception {
        String privateKeyPEM = pemString.replace("-----BEGIN PRIVATE KEY-----", "")
                .replace("-----END PRIVATE KEY-----", "")
                .replaceAll("\\s+", "");
        byte[] encodedPrivateKey = Base64.getMimeDecoder().decode(privateKeyPEM);
        PKCS8EncodedKeySpec keySpec = new PKCS8EncodedKeySpec(encodedPrivateKey);
        KeyFactory keyFactory = KeyFactory.getInstance(ALGORITHM);
        return keyFactory.generatePrivate(keySpec);
    }

    /**
     * 使用公钥进行加密
     */
    public static String encrypt(String plainText, PublicKey publicKey) throws Exception {
        Cipher cipher = Cipher.getInstance(ALGORITHM);
        cipher.init(Cipher.ENCRYPT_MODE, publicKey);
        byte[] encryptedBytes = cipher.doFinal(plainText.getBytes());
        return Base64.getEncoder().encodeToString(encryptedBytes);
    }

    /**
     * 使用私钥进行解密
     */
    public static String decrypt(String encryptedText, PrivateKey privateKey) throws Exception {
        Cipher cipher = Cipher.getInstance(ALGORITHM);
        cipher.init(Cipher.DECRYPT_MODE, privateKey);
        byte[] decryptedBytes = cipher.doFinal(Base64.getDecoder().decode(encryptedText));
        return new String(decryptedBytes);
    }

    /**
     * 保存密钥到文件
     */
    public static void saveKeyToFile(String filePath, String key) throws IOException {
        Files.write(Paths.get(filePath), key.getBytes());
    }

    /**
     * 从文件加载密钥
     */
    public static String loadKeyFromFile(String filePath) throws IOException {
        return new String(Files.readAllBytes(Paths.get(filePath)));
    }

    public static boolean generateKeyPairToFile() {
        try {
            // 生成密钥对
            KeyPair keyPair = generateKeyPair();
            PublicKey publicKey = keyPair.getPublic();
            PrivateKey privateKey = keyPair.getPrivate();

            // 导出密钥为PEM格式
            String publicKeyPEM = exportPublicKeyToPEM(publicKey);
            String privateKeyPEM = exportPrivateKeyToPEM(privateKey);

            System.out.println("Public Key (PEM):");
            System.out.println(publicKeyPEM);
            System.out.println("\nPrivate Key (PEM):");
            System.out.println(privateKeyPEM);

            // 保存密钥到文件
            saveKeyToFile("public_key.pem", publicKeyPEM);
            saveKeyToFile("private_key.pem", privateKeyPEM);
        } catch (Exception e) {
            return false;
        }
        return true;
    }

    public static PublicKey getPublicKey(String filePath) {
        PublicKey loadedPublicKey = null;
        try {
            // 从文件加载密钥
            String loadedPublicKeyPEM = loadKeyFromFile("public_key.pem");
            // 从PEM格式导入密钥
            loadedPublicKey = importPublicKeyFromPEM(loadedPublicKeyPEM);
        } catch (Exception e) {
            e.printStackTrace();
        }
        return loadedPublicKey;
    }

    public static PrivateKey getPrivateKey(String filePath) {
        PrivateKey privateKey = null;
        try {
            // 从文件加载密钥
            String loadedPublicKeyPEM = loadKeyFromFile("private_key.pem");
            // 从PEM格式导入密钥
            privateKey = importPrivateKeyFromPEM(loadedPublicKeyPEM);
        } catch (Exception e) {
            e.printStackTrace();
        }
        return privateKey;
    }


}